import os
import re
import shutil

import numpy as np
import scipy.io as scio
import pandas as pd

from datetime import datetime
from tqdm import tqdm

def process_radardata(allradardatafilepathfolder,allradardatafilepath,chirp):
    """
    function: after "read_mmWave_data(6843).m" processed, this function is save each one's capture start time
    args: allradardatafilepathfolder:the root path of radar data
          allradardatafilepath: one radar data file name
          chirp: the mmWave radar param chirp loops e.g. 64 
    
    return: framefilepath the named with "allradardatafilepath+'_frame.npy'",
            which save radar capture start time 

    e.g.
    a='E:/HARDATASET/RADAR_TRANSFER_STATION/'
    b=os.listdir(a)
    b = sorted(b, key=lambda x: int(re.findall('\d+', x)[0]))
    # print(allradardatafilepathlist)
    for i in tqdm(range(len(b))):
        process_radardata(a,b[i],64)
    """
    Capture_start_time = datetime.strptime(allradardatafilepath, 
                                           '%Y%m%d%H%M%S')

    radardatalen=480000
    onefps=300/(radardatalen/chirp)
    framefilepath=os.path.join(allradardatafilepathfolder,
                               allradardatafilepath,
                               allradardatafilepath+'_frame.npy')


    framelist=[]

    for i in range(int(radardatalen/chirp)):
        framelist.append(i*onefps+Capture_start_time.timestamp())

    framearray=np.array(framelist)
    np.save(framefilepath,framearray)

def process_radar_data(all_action_index_list, SYNC_RADAR_FOLDER='./SYNCHRONIZATION/', LABELED_RADAR_MAT_FOLDER='./TRANSFER_LABELED_RADAR_MAT/'):
    """
    处理毫米波雷达数据，将其切分并保存为MAT文件。

    Args:
        all_action_index_list (dict): 包含时间戳索引的字典。
        SYNC_RADAR_FOLDER (str): 包含毫米波雷达数据的文件夹路径，默认为'./SYNCHRONIZATION/'。
        LABELED_RADAR_MAT_FOLDER (str): 保存切分好的毫米波雷达数据的文件夹路径，默认为'./TRANSFER_LABELED_RADAR_MAT/'。

    Example usage:
    with open('5s_actionIMG_allstartendindex.pkl', 'rb') as f:
        all_action_index_list = pickle.load(f)
    process_radar_data(all_action_index_list, SYNC_RADAR_FOLDER='./SYNCHRONIZATION/', LABELED_RADAR_MAT_FOLDER='./TRANSFER_LABELED_RADAR_MAT/')   
    """
    SYNC_RADAR_FOLDER_LIST = os.listdir(SYNC_RADAR_FOLDER)
    SYNC_RADAR_FOLDER_LIST = sorted(SYNC_RADAR_FOLDER_LIST, key=lambda x: int(re.findall(r'\d+', x)[0]))

    for RADAR_ACTION_TIME in tqdm(SYNC_RADAR_FOLDER_LIST):
        RADAR_START_TIME = os.listdir(os.path.join(SYNC_RADAR_FOLDER, RADAR_ACTION_TIME, 'RADAR'))[0]
        RADAR_ORG_DATAFILEPATH = os.path.join(SYNC_RADAR_FOLDER, RADAR_ACTION_TIME, 'RADAR', RADAR_START_TIME, RADAR_START_TIME + '.mat')
        RADAR_ORG_DATA = scio.loadmat(RADAR_ORG_DATAFILEPATH)  # 解析从MATLAB中处理后得到的数据
        RADAR_DATA = np.moveaxis(np.asarray([RADAR_ORG_DATA['RX1_data'],
                                             RADAR_ORG_DATA['RX2_data'],
                                             RADAR_ORG_DATA['RX3_data'],
                                             RADAR_ORG_DATA['RX4_data']]), 0, -1)

        RADAR_ORG_TIMEFILEPATH = os.path.join(SYNC_RADAR_FOLDER, RADAR_ACTION_TIME, 'RADAR', RADAR_START_TIME, RADAR_START_TIME + '_frame.npy') #读取毫米波雷达时间戳
        RADAR_TIME_DATA = np.load(RADAR_ORG_TIMEFILEPATH)

        VIDEO_TIME_DATA_FILEPATH = os.path.join(SYNC_RADAR_FOLDER, RADAR_ACTION_TIME, 'VIDEO', RADAR_ACTION_TIME, 'frame.pkl') #读取视频文件时间戳
        VIDEO_TIME_DATA = np.load(VIDEO_TIME_DATA_FILEPATH, allow_pickle=True)

        time_action_list = all_action_index_list[RADAR_ACTION_TIME] #读取每隔5s的完整时间戳

        LABELED_RADAR_MAT_time = os.path.join(LABELED_RADAR_MAT_FOLDER, RADAR_ACTION_TIME)

        if not os.path.exists(LABELED_RADAR_MAT_time):
            os.mkdir(LABELED_RADAR_MAT_time)

        for action_name_index, action_index_list in time_action_list.items():

            for i in range(len(action_index_list)):
                video_start_timestamp = VIDEO_TIME_DATA[action_index_list[i][0]].timestamp()
                video_end_timestamp = VIDEO_TIME_DATA[action_index_list[i][1]].timestamp()
                RADAR_ACTION_STARTTIME_LIST = []
                RADAR_ACTION_ENDTIME_LIST = []
                #获取毫米波雷达和视频数据对应的时间戳
                for radar_time in RADAR_TIME_DATA:
                    if round(video_start_timestamp, 1) == round(radar_time, 1):
                        RADAR_ACTION_STARTTIME_LIST.append(radar_time)

                    if round(video_end_timestamp, 1) == round(radar_time, 1):
                        RADAR_ACTION_ENDTIME_LIST.append(radar_time)

                RADAR_ACTION_STARTTIME = RADAR_ACTION_STARTTIME_LIST[0]
                RADAR_ACTION_ENDTIME = RADAR_ACTION_ENDTIME_LIST[0]

                RADAR_ACTION_STARTFPS = round((RADAR_ACTION_STARTTIME - RADAR_TIME_DATA[0]) * 25)
                RADAR_ACTION_ENDFPS = RADAR_ACTION_STARTFPS + 100
                ONEACTION_RADAR_ORG_DATA = RADAR_DATA[:, RADAR_ACTION_STARTFPS * 64:RADAR_ACTION_ENDFPS * 64, :]

                RADAR_ACTION_MAT_FILEPATH = os.path.join(LABELED_RADAR_MAT_time, str(action_name_index) + '_' + str(i) + '.mat')

                scio.savemat(RADAR_ACTION_MAT_FILEPATH, {'RADAR_DATA': ONEACTION_RADAR_ORG_DATA})

